package collections;

import util.Tracer;

public class RedBlackTree<T extends Comparable<T>> {
	private RedBlackTreeNode<T> root;
	
	public RedBlackTreeNode<T> search(T k)
	{
		RedBlackTreeNode<T> node = root;

		int ret;
		while(node != null && (ret = k.compareTo(node.getKey())) != 0)
		{
			if(ret < 0)
				node = node.getLeftChild();
			else
				node = node.getRightChild();
		}
		
		return node;
	}
	
	public void insert(RedBlackTreeNode<T> node)
	{
		RedBlackTreeNode<T> parent = null;
		RedBlackTreeNode<T> tmp = root;
		
		while(tmp != null)
		{
			parent = tmp;

			if(node.getKey().compareTo(tmp.getKey()) < 0)
			{
				tmp = tmp.getLeftChild();
			}
			else
			{
				tmp = tmp.getRightChild();
			}
		}
		
		if(parent == null)
		{
			setRoot(node);
		}
		else if(node.getKey().compareTo(parent.getKey()) < 0)
		{
			parent.setLeftChild(node);
		}
		else
		{
			parent.setRightChild(node);
		}
		
		node.setColor(RedBlackTreeColor.RED);
		insertionFixup(node);
	}
	
	private void insertionFixup(RedBlackTreeNode<T> node)
	{
		RedBlackTreeNode<T> x = node;
		RedBlackTreeNode<T> parent;
		
		while(x != root && (parent = x.getParent()).getColor() == RedBlackTreeColor.RED)
		{
			RedBlackTreeNode<T> grandParent = parent.getParent();
			if(parent == grandParent.getLeftChild())
			{
				RedBlackTreeNode<T> uncle = grandParent.getRightChild();
				if(uncle != null && uncle.getColor() == RedBlackTreeColor.RED)
				{
					grandParent.setColor(RedBlackTreeColor.RED);
					parent.setColor(RedBlackTreeColor.BLACK);
					uncle.setColor(RedBlackTreeColor.BLACK);
					x = grandParent;
				}
				else
				{
					if(x == parent.getRightChild())
					{
						leftRotate(parent);
						
						RedBlackTreeNode<T> tmp = parent;
						parent = x;
						x = tmp;
					}
					
					parent.setColor(RedBlackTreeColor.BLACK);
					grandParent.setColor(RedBlackTreeColor.RED);
					rightRotate(grandParent);
				}
			}
			else
			{
				RedBlackTreeNode<T> uncle = grandParent.getLeftChild();
				if(uncle != null && uncle.getColor() == RedBlackTreeColor.RED)
				{
					grandParent.setColor(RedBlackTreeColor.RED);
					parent.setColor(RedBlackTreeColor.BLACK);
					uncle.setColor(RedBlackTreeColor.BLACK);
					x = grandParent;
				}
				else
				{
					if(x == parent.getLeftChild())
					{
						rightRotate(parent);
						
						RedBlackTreeNode<T> tmp = parent;
						parent = x;
						x = tmp;
					}
					
					grandParent.setColor(RedBlackTreeColor.RED);
					parent.setColor(RedBlackTreeColor.BLACK);
					leftRotate(grandParent);
				}
			}
		}
		
		root.setColor(RedBlackTreeColor.BLACK);
	}
	
	public void remove(RedBlackTreeNode<T> node)
	{
		if(node.getLeftChild() == null)
		{
			transplant(node, node.getRightChild());
		}
		else if(node.getRightChild() == null)
		{
			transplant(node, node.getLeftChild());
		}
		else
		{
			RedBlackTreeNode<T> replacer = minimum(node.getRightChild());
			
			if(replacer.getParent() != node)
			{
				transplant(replacer, replacer.getRightChild());
				replacer.setRightChild(node.getRightChild());
			}
			
			transplant(node, replacer);
			replacer.setLeftChild(node.getLeftChild());
		}
	}
	
	public RedBlackTreeNode<T> minimum()
	{
		return minimum(root);
	}
	
	private RedBlackTreeNode<T> minimum(RedBlackTreeNode<T> node)
	{
		if(node == null)
			return null;

		RedBlackTreeNode<T> child;
		while((child = node.getLeftChild()) != null)
		{
			node = child;
		}
		
		return node;
	}
	
	private void transplant(RedBlackTreeNode<T> node, RedBlackTreeNode<T> replacer)
	{
		RedBlackTreeNode<T> parent = node.getParent();
		if(parent == null)
			setRoot(replacer);
		else if(node == parent.getLeftChild())
			parent.setLeftChild(replacer);
		else if(node == parent.getRightChild())
			parent.setRightChild(replacer);
	}
	
	private void leftRotate(RedBlackTreeNode<T> node)
	{
		RedBlackTreeNode<T> rightChild = node.getRightChild();
		node.setRightChild(rightChild.getLeftChild());

		RedBlackTreeNode<T> parent = node.getParent();
		if(parent == null)
			setRoot(rightChild);
		else if(node == parent.getLeftChild())
			parent.setLeftChild(rightChild);
		else
			parent.setRightChild(rightChild);
		
		rightChild.setLeftChild(node);
	}
	
	private void rightRotate(RedBlackTreeNode<T> node)
	{
		RedBlackTreeNode<T> leftChild = node.getLeftChild();
		node.setLeftChild(leftChild.getRightChild());

		RedBlackTreeNode<T> parent = node.getParent();
		if(parent == null)
			setRoot(leftChild);
		else if(node == parent.getLeftChild())
			parent.setLeftChild(leftChild);
		else
			parent.setRightChild(leftChild);
		
		leftChild.setRightChild(node);
	}
	
	private void inorderWalk(RedBlackTreeNode<T> node)
	{
		if(node == null) return;
		inorderWalk(node.getLeftChild());
		
		RedBlackTreeNode<T> parent = node.getParent();
		Tracer.trace(parent == null ? null : parent.getKey(), node.getKey(), node.getColor());
		inorderWalk(node.getRightChild());
	}
	
	private void setRoot(RedBlackTreeNode<T> node)
	{
		node.setParent(null);
		root = node;
	}
	
	public static void main(String[] args) {
		RedBlackTree<Integer> tree = new RedBlackTree<>();
		tree.insert(new RedBlackTreeNode<Integer>(6));
		tree.insert(new RedBlackTreeNode<Integer>(7));
		tree.insert(new RedBlackTreeNode<Integer>(8));
		tree.insert(new RedBlackTreeNode<Integer>(9));
		tree.insert(new RedBlackTreeNode<Integer>(10));
		tree.insert(new RedBlackTreeNode<Integer>(11));
		tree.insert(new RedBlackTreeNode<Integer>(12));
		tree.insert(new RedBlackTreeNode<Integer>(1));
		tree.insert(new RedBlackTreeNode<Integer>(2));
		tree.insert(new RedBlackTreeNode<Integer>(3));
		tree.insert(new RedBlackTreeNode<Integer>(4));
		tree.insert(new RedBlackTreeNode<Integer>(5));
		tree.inorderWalk(tree.root);
	}
}
